#! /usr/bin/python
import sys
import pefile

if len(sys.argv) < 3:
    print("[*] Usage: <target DLL file path> <optional: --export-name [export name to force delta from]> [Gadget RVA(s)/VAs]")
    exit()

TargetFilePath=sys.argv[1]
OverrideExportName = None

print("... opening file at " + TargetFilePath + " to hunt for nearest exports to specified gadgets")

pe = pefile.PE(TargetFilePath)

if not hasattr(pe, 'DIRECTORY_ENTRY_EXPORT'):
    print("... EAT missing. Invalid target DLL " + TargetFilePath)
    sys.exit(-1)
   
GadgetRvaList = []
GadgetCount = 0
GadgetsIndex = 2
   
if sys.argv[2] == "--export-name":
    OverrideExportName = sys.argv[3]
    GadgetsIndex = 4
           

for ArgIndex in range(GadgetsIndex, len(sys.argv)):
    CurrentGadgetAddress = int(sys.argv[ArgIndex][2:], 16)

    if CurrentGadgetAddress > pe.OPTIONAL_HEADER.ImageBase:
        if CurrentGadgetAddress < (pe.OPTIONAL_HEADER.ImageBase + pe.OPTIONAL_HEADER.SizeOfImage):
            CurrentGadgetAddress = (CurrentGadgetAddress - pe.OPTIONAL_HEADER.ImageBase)
            
    GadgetRvaList.append(CurrentGadgetAddress)
    GadgetCount += 1
    #ArgIndex += 1

#print("... target gadget list:")

#for GadgetRva in GadgetRvaList:
#    print("0x%08x" % (GadgetRva))
    
if OverrideExportName != None:
    print("... hunting for delta only relative to export %s"%(OverrideExportName))
    
ExportDict = {}
ExportRvaList = []

# Make a list of all export RVAs as well as a dictionary with RVAs as keys and names as values.
# While buiilding the list/dictionary, find the export which is LOWER (or equal) than the target gadget but as close as possible

for ExportEntry in pe.DIRECTORY_ENTRY_EXPORT.symbols:
    if ExportEntry.name:
        if OverrideExportName != None:
            if OverrideExportName == ExportEntry.name.decode('utf-8'):
                ExportDict[ExportEntry.address] = ExportEntry.name.decode('utf-8')
                ExportRvaList.append(ExportEntry.address)
                break
        else:
            ExportDict[ExportEntry.address] = ExportEntry.name.decode('utf-8')
            ExportRvaList.append(ExportEntry.address)
            
        #print("%s - 0x%08x" % (ExportEntry.name.decode('utf-8'), ExportEntry.address))
    else:
        print("... invalid export name")

print("... generating export dictionary with %d entries"%(len(ExportDict)))

#for ExportEntry in pe.DIRECTORY_ENTRY_EXPORT.symbols:
#    if ExportEntry.ordinal:
#        print("... export ordinal: 0x%08x" % (ExportEntry.ordinal))
#    else:
#        print("... no ordinal")
    
#print("... export RVA list:")

#for ExportRva in ExportRvaList:
#    print("0x%08x"%(ExportRva))
   
ExportRvaList.sort()   
#print("... sorted export RVA list:")

#for ExportRva in ExportRvaList:
#    print("0x%08x"%(ExportRva))

print("... calculating nearest exports with forward deltas...")

for GadgetRva in GadgetRvaList:
    NearestRva = -1
    NearestDelta = -1
    NearestRvaIndex = -1
    CurrentRvaIndex = 0

    for ExportRva in ExportRvaList:
        CurrentDelta = -1
        
        if GadgetRva >= ExportRva:
            CurrentDelta = (GadgetRva - ExportRva)
             
        if NearestRva == -1:
            if CurrentDelta >= 0:
                NearestRva = ExportRva
                NearestDelta = CurrentDelta
                NearestRvaIndex = CurrentRvaIndex
        else:
            if CurrentDelta >= 0:
                if CurrentDelta < NearestDelta:
                    NearestRva = ExportRva
                    NearestDelta = CurrentDelta
                    NearestRvaIndex = CurrentRvaIndex
                    
        CurrentRvaIndex += 1

    if NearestRva == -1:
        print("... failed to identify any exports which came prior to gadget at 0x%08x" %(GadgetRva))
    else:
        if (NearestRvaIndex + 1) < len(ExportRvaList):
            print("%s:0x%08x (+0x%x) <- 0x%08x -> (+0x%x) %s:0x%08x" % (ExportDict[ExportRvaList[NearestRvaIndex]], ExportRvaList[NearestRvaIndex], NearestDelta, GadgetRva, ExportRvaList[NearestRvaIndex + 1] - GadgetRva, ExportDict[ExportRvaList[NearestRvaIndex + 1]], ExportRvaList[NearestRvaIndex + 1]))
        else:
            print("%s:0x%08x <- 0x%08x -> end of .text" % (ExportDict[ExportRvaList[NearestRvaIndex]], ExportRvaList[NearestRvaIndex], GadgetRva))
        
print("... calculating nearest exports with backward deltas...")

for GadgetRva in GadgetRvaList:
    NearestRva = 0
    NearestDelta = 0
    CurrentRvaIndex = 0

    for ExportRva in ExportRvaList:
        CurrentDelta = 0
        
        if GadgetRva < ExportRva:
            CurrentDelta = (GadgetRva - ExportRva)
             
        if NearestRva == 0:
            if CurrentDelta < 0:
                NearestRva = ExportRva
                NearestDelta = CurrentDelta
                NearestRvaIndex = CurrentRvaIndex
        else:
            if CurrentDelta < 0:
                if CurrentDelta > NearestDelta:
                    NearestRva = ExportRva
                    NearestDelta = CurrentDelta
                    NearestRvaIndex = CurrentRvaIndex
                    
        CurrentRvaIndex += 1

    if NearestRva == 0:
        print("... failed to identify any exports which came prior to gadget at 0x%08x" %(GadgetRva))
    else:
        if (NearestRvaIndex - 1) >= 0:
            print("%s:0x%08x (+0x%x) <- 0x%08x -> (+0x%x) %s:0x%08x" % (ExportDict[ExportRvaList[NearestRvaIndex - 1]], ExportRvaList[NearestRvaIndex - 1],  GadgetRva - ExportRvaList[NearestRvaIndex - 1], GadgetRva, ExportRvaList[NearestRvaIndex] - GadgetRva, ExportDict[ExportRvaList[NearestRvaIndex]], ExportRvaList[NearestRvaIndex]))
        else:
            print("start of .text <- 0x%08x -> %s:0x%08x" % (GadgetRva, ExportDict[ExportRvaList[NearestRvaIndex]], ExportRvaList[NearestRvaIndex]))